package main

import (
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"fmt"
	"github.com/Unixeno/TheMoon/authentication"
	"io"
	"log"
	"math/big"
	"os"
	"time"
)

var (
	validFor = flag.Duration("duration", 365*24*time.Hour, "证书的有效时间，单位天")
	isCA     = flag.Bool("ca", false, "是否生成CA证书")
	keyType  = flag.String("type", "ed25519", "生成密钥对时使用的算法类型，支持rsa2048, rsa4096, ed25519, ecdsaP224, ecdsaP256, ecdsaP384, ecdsaP521")
	certOut  = flag.String("out", "", "输出证书路径")
	keyOut   = flag.String("keyout", "", "输出私钥路径")
	caFile   = flag.String("cain", "", "使用该ca证书生成证书")
	caKeyFile = flag.String("cakey", "", "生成证书时使用的ca密钥")
	node     = flag.String("node", "", "node ID")
)

func main() {
	flag.Parse()

	certOutFile := os.Stdout
	keyOutFile := os.Stdout
	var err error

	if len(*certOut) != 0 {
		if certOutFile, err = os.OpenFile(*certOut, os.O_RDWR|os.O_CREATE, 0644); err != nil {
			log.Fatalf("无法打开证书文件: %v", err)
		}
	}
	defer certOutFile.Close()

	if len(*keyOut) != 0 {
		if keyOutFile, err = os.OpenFile(*keyOut, os.O_RDWR|os.O_CREATE, 0644); err != nil {
			log.Fatalf("无法打开私钥文件: %v", err)
		}
	}
	defer keyOutFile.Close()

	log.Printf("使用 %s 生成密钥对", *keyType)

	privateKey, err := genPrivateKey(*keyType)
	if err != nil {
		log.Fatalf("failed to generate private key: %v", err)
	}

	log.Println("生成证书...")

	if *isCA {
		err := genCA(*validFor, privateKey, certOutFile, keyOutFile)
		if err != nil {
			log.Fatalf("CA证书生成失败: %v", err)
		}
	} else {
		if len(*caFile) == 0 || len(*caKeyFile) == 0 {
			log.Fatalf("生成普通证书时必须指定ca证书和ca密钥")
		}
		ca, err := authentication.ReadCertificate(*caFile)
		if err != nil {
			log.Fatalf("读取ca证书文件时发生错误: %v", err)
		}
		caKey, err := authentication.ReadKeyFile(*caKeyFile)
		if err != nil {
			log.Fatalf("读取ca私钥时发生错误: %v", err)
		}

		err = genCert(*validFor, *node, ca, caKey, privateKey, certOutFile, keyOutFile)
		if err != nil {
			log.Fatalf("证书生成失败: %v", err)
		} else {
			fmt.Printf("证书: %s\n私钥: %s\n\n", *certOut, *keyOut)
		}
	}
}

func genSerialNumber() *big.Int {
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		log.Fatalf("Failed to generate serial number: %v", err)
	}
	return serialNumber
}

func genCA(expire time.Duration, caPrivateKey interface{}, certOut io.Writer, keyOut io.Writer) error {
	ca := &x509.Certificate{
		SerialNumber: genSerialNumber(),
		Subject: pkix.Name{
			Organization:  []string{"Islands Developers"},
			Country:       []string{"CN"},
			Province:      []string{"Shanxi"},
			CommonName:    "Planet Test CA",
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(expire),
		IsCA:                  true,
		//ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		BasicConstraintsValid: true,
	}

	caBytes, err := x509.CreateCertificate(rand.Reader, ca, ca, publicKey(caPrivateKey), caPrivateKey)
	if err != nil {
		return err
	}
	return dumpCert(caBytes, certOut, caPrivateKey, keyOut)
}

func genCert(expire time.Duration, node string, parent *x509.Certificate, parentKey interface{}, privateKey interface{}, certOut io.Writer, keyOut io.Writer) error {
	cert := &x509.Certificate{
		SerialNumber: genSerialNumber(),
		Subject: pkix.Name{
			Organization:  []string{"Islands Developers"},
			Country:       []string{"CN"},
			Province:      []string{"Shanxi"},
			CommonName:    node,
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(expire),
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageDataEncipherment,
		BasicConstraintsValid: true,
	}
	certBytes, err := x509.CreateCertificate(rand.Reader, cert, parent, publicKey(privateKey), parentKey)
	if err != nil {
		return err
	}

	signedCertificate, _ := x509.ParseCertificate(certBytes)
	printCertificate(signedCertificate)

	return dumpCert(certBytes, certOut, privateKey, keyOut)
}

func dumpCert(cert []byte, certOut io.Writer, privateKey interface{}, keyOut io.Writer) error {
	if err := pem.Encode(certOut, &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: cert,
	}); err != nil {
		return err
	}

	if keyBytes, err := x509.MarshalPKCS8PrivateKey(privateKey); err != nil {
		return err
	} else {
		if err := pem.Encode(keyOut, &pem.Block{
			Type:  "PRIVATE KEY",
			Bytes: keyBytes,
		}); err != nil {
			return err
		}
	}
	return nil
}

func genPrivateKey(keyType string) (priv interface{}, err error) {
	switch keyType {
	case "ed25519":
		_, priv, err = ed25519.GenerateKey(rand.Reader)
	case "rsa2048":
		priv, err = rsa.GenerateKey(rand.Reader, 2048)
	case "rsa4096":
		priv, err = rsa.GenerateKey(rand.Reader, 4096)
	case "ecdsaP224":
		priv, err = ecdsa.GenerateKey(elliptic.P224(), rand.Reader)
	case "ecdsaP256":
		priv, err = ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	case "ecdsaP384":
		priv, err = ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	case "ecdsaP521":
		priv, err = ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
	default:
		err = fmt.Errorf("无效的key类型 %s", keyType)
	}
	return
}

func publicKey(privateKey interface{}) interface{} {
	switch k := privateKey.(type) {
	case *rsa.PrivateKey:
		return &k.PublicKey
	case *ecdsa.PrivateKey:
		return &k.PublicKey
	case ed25519.PrivateKey:
		return k.Public().(ed25519.PublicKey)
	default:
		return nil
	}
}

func printCertificate(cert *x509.Certificate)  {
	fmt.Printf("\n=======Certificate======\n")
	fmt.Printf("序列号: %X\n", cert.SerialNumber)
	fmt.Printf("SubjectCN: %s\n", cert.Subject.CommonName)
	fmt.Printf("有效期: %s\n", cert.NotAfter.Format("2006-01-02T15:04:05"))
	fmt.Printf("Signature: %X\n", cert.Signature)
}
